package com.zyq.intercepter;

import cn.hutool.core.util.StrUtil;
import com.zyq.annotation.AccessLimit;
import com.zyq.constant.AccessLimits;
import com.zyq.enuns.ResponseCode;
import com.zyq.exception.BusinessException;
import com.zyq.util.RedisUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.stereotype.Component;
import org.springframework.web.method.HandlerMethod;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.net.InetAddress;
import java.net.UnknownHostException;
import java.time.LocalDateTime;
import java.util.Collections;
import java.util.List;
import java.util.Objects;
import java.util.concurrent.TimeUnit;

/**
 * Redis限流拦截器
 * @author xiaochi
 * @date 2021/12/11 22:13
 */
@Slf4j
@Component("accessLimitIntercepter")
public class AccessLimitIntercepter extends AbstractBaseInterceptor{

    @Override
    protected boolean execute(HttpServletRequest request, HttpServletResponse response, HandlerMethod method) {
        try {
            AccessLimit accessLimit = method.getMethodAnnotation(AccessLimit.class);
            String key = getKey(request,accessLimit,method);
            AccessLimits.StoreType storeType = accessLimit.store();
            // redis
            if (Objects.equals(storeType, AccessLimits.StoreType.REDIS)){
                int limitCount = accessLimit.count();
                key = accessLimit.prefix() + key;
                // 限流
                if (Objects.equals(accessLimit.type(), AccessLimits.Type.LIMIT)){
                    return limitRedis(key, Collections.singletonList(key),limitCount,accessLimit.timeOut()==0?600:accessLimit.timeOut(),accessLimit.timeOutUnit());
                }
                // 限制访问次数
                if (Objects.equals(accessLimit.type(),AccessLimits.Type.COUNT)){
                    //ImmutableList<String> keys = ImmutableList.of(StringUtils.join(accessLimit.prefix(),key));
                    int timeOut = accessLimit.timeOut();
                    return limitRedisCount(key,Collections.singletonList(key),limitCount,Objects.equals(timeOut,0)?(3600*24*1000):timeOut,accessLimit.timeOutUnit());
                }
            }
        } catch (Throwable e) {
            log.error("accessLimitIntercepter限流:{}",e);
            if (e.getCause() instanceof BusinessException){
                throw (BusinessException) e.getCause();
            }
            throw new RuntimeException(e);
        }
        return false;
    }

    /**
     * redis 限流
     * @param keys
     * @param limitCount
     * @param timeOut
     * @param timeUnit
     * @return
     */
    private boolean limitRedis(String key, List<String> keys, int limitCount, int timeOut, TimeUnit timeUnit) throws Throwable {
        /*Object count = RedisUtil.instance.getRedisTemplate().opsForValue().get(key);
        if (null == count) {
            RedisUtil.instance.getRedisTemplate().opsForValue().set(key,1,timeOut,timeUnit);
            return null;
        }
        if (Integer.parseInt(String.valueOf(count)) >= limitCount){
            throw new BusinessException("服务器繁忙，请稍后再试!");
        }
        RedisUtil.instance.getRedisTemplate().opsForValue().increment(key);
        return null;*/
        String luaScript = buildLuaScript();
        RedisScript<Number> redisScript = new DefaultRedisScript<>(luaScript, Number.class);
        // redisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource("concurrent_request_rate_limiter.lua")));
        Number count = RedisUtil.instance.getRedisTemplate().execute(redisScript, keys, String.valueOf(limitCount), String.valueOf(timeOut / 1000),timeUnit);// timeOut/1000 是因为默认为秒
        //logger.info("Access try count is {} and key = {}", count, key);
        if(count == null || count.intValue() > limitCount){
            log.warn("<<================= redis 限流,key:{},每秒只能count:{}次请求,当前时间:{}", key,limitCount, LocalDateTime.now());
            throw new BusinessException(ResponseCode.SYSTEM_BUSY_AGAIN_LATER);
        }
        return true;
    }

    /**
     * redis 限制访问次数
     * @param key
     * @param keys
     * @param limitCount
     * @param timeOut
     * @return
     */
    private boolean limitRedisCount(String key, List<String> keys, int limitCount, int timeOut,TimeUnit timeUnit){
        String luaScript = buildLuaScript();
        RedisScript<Number> redisScript = new DefaultRedisScript<>(luaScript, Number.class);
        Number count = RedisUtil.instance.getRedisTemplate().execute(redisScript, keys, String.valueOf(limitCount), String.valueOf(timeOut / 1000),timeUnit);// timeOut/1000 是因为默认为秒
        //logger.info("Access try count is {} and key = {}", count, key);
        if(count == null || count.intValue() > limitCount){
            log.warn("<<================= 访问超限,key:{},count:{},当前时间:{}", key,limitCount,LocalDateTime.now());
            throw new BusinessException(ResponseCode.ACCESS_NUMBER_LIMIT);
        }
        return true;
    }

    /**
     * lua限流脚本
     * @return
     */
    public String buildLuaScript(){
        StringBuilder sb = new StringBuilder();
        //定义c
        sb.append("local c");
        //获取redis中的值
        sb.append("\nc = redis.call('get',KEYS[1])");
        //如果调用不超过最大值
        sb.append("\nif c and tonumber(c) > tonumber(ARGV[1]) then");
        //直接返回
        sb.append("\n return c;");
        //结束
        sb.append("\nend");
        //访问次数加一
        sb.append("\nc = redis.call('incr',KEYS[1])");
        //如果是第一次调用
        sb.append("\nif tonumber(c) == 1 then");
        //设置对应值的过期设置
        sb.append("\nredis.call('expire',KEYS[1],ARGV[2])");
        //结束
        sb.append("\nend");
        //返回
        sb.append("\nreturn c;");
        return sb.toString();
    }

    /**
     * 获取key
     * @param request
     * @param accessLimit
     * @param handlerMethod
     * @return
     */
    private String getKey(HttpServletRequest request, AccessLimit accessLimit, HandlerMethod handlerMethod){
        AccessLimits.ResourceType limitResourceType = accessLimit.resourceType();
        //如果类型是IP，则根据IP限制访问次数，key取IP地址
        if (AccessLimits.ResourceType.IP.equals(limitResourceType)){
            return getIpAddr(request) + accessLimit.separator() + request.getRequestURI();
        }
        //如果类型是customer，则根据key限制访问次数
        if (AccessLimits.ResourceType.CUSTOMER.equals(limitResourceType)){
            return this.getSpelValueByKey(accessLimit.key(), request);
        }
        //否则按照方法名称限制访问次数
        return handlerMethod.getMethod().getName();
    }

    private static final String IP_UTILS_FLAG = ",";
    private static final String UNKNOWN = "unknown";
    private static final String LOCALHOST_IP = "0:0:0:0:0:0:0:1";
    private static final String LOCALHOST_IP1 = "127.0.0.1";
    /**
     * 获取IP地址
     * <p>
     * 使用Nginx等反向代理软件， 则不能通过request.getRemoteAddr()获取IP地址
     * 如果使用了多级反向代理的话，X-Forwarded-For的值并不止一个，而是一串IP地址，X-Forwarded-For中第一个非unknown的有效IP字符串，则为真实IP地址
     */
    public String getIpAddr(HttpServletRequest request) {
        String ip = null;
        try {
            //以下两个获取在k8s中，将真实的客户端IP，放到了x-Original-Forwarded-For。而将WAF的回源地址放到了 x-Forwarded-For了。
            ip = request.getHeader("X-Original-Forwarded-For");
            if (StrUtil.isEmpty(ip) || UNKNOWN.equalsIgnoreCase(ip)) {
                ip = request.getHeader("X-Forwarded-For");
            }
            //获取nginx等代理的ip
            if (StrUtil.isEmpty(ip) || UNKNOWN.equalsIgnoreCase(ip)) {
                ip = request.getHeader("x-forwarded-for");
            }
            if (StrUtil.isEmpty(ip) || UNKNOWN.equalsIgnoreCase(ip)) {
                ip = request.getHeader("Proxy-Client-IP");
            }
            if (StrUtil.isEmpty(ip) || ip.length() == 0 || UNKNOWN.equalsIgnoreCase(ip)) {
                ip = request.getHeader("WL-Proxy-Client-IP");
            }
            if (StrUtil.isEmpty(ip) || UNKNOWN.equalsIgnoreCase(ip)) {
                ip = request.getHeader("HTTP_CLIENT_IP");
            }
            if (StrUtil.isEmpty(ip) || UNKNOWN.equalsIgnoreCase(ip)) {
                ip = request.getHeader("HTTP_X_FORWARDED_FOR");
            }
            //兼容k8s集群获取ip
            if (StrUtil.isEmpty(ip) || UNKNOWN.equalsIgnoreCase(ip)) {
                ip = request.getRemoteAddr();
                if (LOCALHOST_IP1.equalsIgnoreCase(ip) || LOCALHOST_IP.equalsIgnoreCase(ip)) {
                    //根据网卡取本机配置的IP
                    InetAddress iNet = null;
                    try {
                        iNet = InetAddress.getLocalHost();
                    } catch (UnknownHostException e) {
                        log.error("getClientIp error: {}", e);
                    }
                    ip = iNet.getHostAddress();
                }
            }
        } catch (Exception e) {
            log.error("IPUtils ERROR ", e);
        }
        //使用代理，则获取第一个IP地址
        if (!StrUtil.isEmpty(ip) && ip.indexOf(IP_UTILS_FLAG) > 0) {
            ip = ip.substring(0, ip.indexOf(IP_UTILS_FLAG));
        }
        return ip;
    }
}
